from __future__ import absolute_import, division, print_function, unicode_literals

import datetime  # For datetime objects
import os.path  # To manage paths
import sys  # To find out the script name (in argv[0])

# Import the backtrader platform
import backtrader as bt

import MySQLDataFeed as mdf
import SimpleMySqlClass as smc
import talib as talib
import numpy as np
import akshare as ak
import pandas as pd

# 1，输入code
# 2，检查沪深300，中正500是否在内
# 3，检查 t_stock_data_daily_${code} 是否建表，创建表
# 4，拉取code日线数据，落库
# 5，跑双均线模型回测
# 6，记录开仓买卖记录到流水表


def log(out_put):
    print("%s", out_put)


def get_cfg_info(conn, code):
    sql_300 = f"select * from t_cfg_hs_300 where code='{code}' limit 1"
    cfg_300_obj = conn.get_one(sql_300)
    print("cfg_300_obj", cfg_300_obj)
    if cfg_300_obj:
        print(f"code=%{code},是沪深300成分股")
        return cfg_300_obj

    sql_500 = f"select * from t_cfg_zz_500 where code='{code}' limit 1"
    cfg_500_obj = conn.get_one(sql_500)
    if cfg_500_obj:
        print(f"code=%{code},是中证500成分股")
        return cfg_500_obj

    return None


# 是否是沪深300，中证500
def check_code_hs300_zz500(conn, code):
    condition = f" code = '%{code}'"
    check_hs_300_rs = conn.count("t_cfg_hs_300", condition)
    check_zz_500_rs = conn.count("t_cfg_zz_500", condition)
    return check_hs_300_rs + check_zz_500_rs > 0


def check_and_create_table(conn, code):
    base_table_name = "t_trade_data"
    table_name = base_table_name + "_" + code
    if not check_table_name(conn, table_name):
        create_like_table(conn, base_table_name, table_name)
    return


def check_table_name(conn, table_name):
    all_tables = conn.check_table(table_name)
    for t in all_tables:
        cur_table_name = t[0]
        print("table = ", cur_table_name)
        if cur_table_name == table_name:
            return True
    return False


def create_like_table(conn, base_table_name, table_name):
    copy_sql = f"create table {table_name} like {base_table_name}"
    rs = conn.execute(copy_sql)
    return rs


# 2, save data to mysql
def save_data_to_mysql(conn, stock_data_df):
    rs_ct = 0
    for index, row in stock_data_df.iterrows():
        # 插入新记录
        data = {
            "id": 0,
            "date": row["date"],
            "date_time": row["date_time"],
            "code": row["code"],
            "open": row["open"],
            "high": row["high"],
            "low": row["low"],
            "close": row["close"],
            "volume": row["volume"],
            "amount": row["amount"],
            "create_time": datetime.datetime.now(),
            "update_time": datetime.datetime.now(),
        }
        table_name = "t_trade_data_" + row["code"]
        rs = conn.insertIgnore(table_name, data)
        rs_ct = rs_ct + rs

    print("------------------------rs_ct=", rs_ct)
    return


# 3, query data from mysql
def read_data_from_mysql(conn, table_name):
    rs = conn.count(table_name, None)
    print("query rs = ", rs)
    pass


# 获取数据库链接
def get_mysql_conn():
    conn = smc.SimpleMySqlClass("127.0.0.1", "stocks", "haojg", "haojg")
    return conn


# 获取最近15个月的数据
def get_start_date_yyyyMMdd():
    end_date = datetime.datetime.now()
    # 一年+3个月
    start_date = end_date - datetime.timedelta(days=365 + 92)
    start_date_param = start_date.strftime("%Y%m%d")
    return start_date_param


# 当前时间为结束日期
def get_end_date_yyyyMMdd():
    end_date = datetime.datetime.now()
    end_date_param = end_date.strftime("%Y%m%d")
    return end_date_param


def get_daily_data(code):
    start_date_param = get_start_date_yyyyMMdd()
    end_date_param = get_end_date_yyyyMMdd()

    stock_qfq_data = ak.stock_zh_a_hist(
        symbol=code, adjust="qfq", start_date=start_date_param, end_date=end_date_param
    )
    # print(stock_qfq_data)
    stock_qfq_data = stock_qfq_data.iloc[:, 0:7]

    # 日期     开盘     收盘     最高     最低     成交量  成交额
    # 列名记得这样定义好
    stock_qfq_data.columns = [
        "date",
        "open",
        "close",
        "high",
        "low",
        "volume",
        "amount",
    ]
    stock_qfq_data["date_time"] = stock_qfq_data.loc[:, ["date"]]
    stock_qfq_data["code"] = code
    # print(stock_qfq_data)
    return stock_qfq_data


def get_data_and_read(code):
    # 获取链接
    conn = get_mysql_conn()
    cfg_obj = get_cfg_info(conn, code)
    print("检查是否成分股=", cfg_obj)
    table_name = "t_trade_data_" + code
    code_desc = cfg_obj[3] if cfg_obj else "-"
    print(f"开始获取{ code_desc} 数据, code = {code}, table_name={table_name}")
    check_and_create_table(conn, code)
    # 拉取远端数据
    stock_qfq_data = get_daily_data(code)
    # 落库
    save_data_to_mysql(conn, stock_qfq_data)

    # 读数
    read_data_from_mysql(conn, table_name)
    # 关闭连接
    conn.close()


def backtrader_test_by_data(code):
    # Create a cerebro entity
    cerebro = bt.Cerebro()

    # 5.加载观察者 Observers
    # cerebro.addobserver(OrderObserver)

    # Add a strategy
    cerebro.addstrategy(MyMaMacdRsi)
    # add kdj 指标
    # cerebro.addindicator(KDJ)

    tableName = "t_trade_data_" + code
    sql = f"select `date`, `open`, `high`, `low`, `close`, `volume` from {tableName}"

    data = mdf.MySQLDataFeed(sql)

    # Add the Data Feed to Cerebro
    cerebro.adddata(data, name="2Line")

    # Set our desired cash start
    # 初始资金 100,000,000
    cerebro.broker.setcash(100000.0)
    # 佣金，双边各 0.0003
    cerebro.broker.setcommission(commission=0.0003)
    # 滑点：双边各 0.0001
    cerebro.broker.set_slippage_perc(perc=0.0001)

    # Add a FixedSize sizer according to the stake
    cerebro.addsizer(bt.sizers.FixedSize, stake=100)

    # Set the commission

    cerebro.addanalyzer(bt.analyzers.TimeReturn, _name="pnl")  # 返回收益率时序数据
    # 添加分析指标
    # 返回年初至年末的年度收益率
    cerebro.addanalyzer(bt.analyzers.AnnualReturn, _name="_AnnualReturn")
    # 计算最大回撤相关指标
    cerebro.addanalyzer(bt.analyzers.DrawDown, _name="_DrawDown")
    # 计算年化收益：日度收益
    cerebro.addanalyzer(bt.analyzers.Returns, _name="_Returns", tann=252)
    # 计算年化夏普比率：日度收益
    cerebro.addanalyzer(
        bt.analyzers.SharpeRatio,
        _name="_SharpeRatio",
        timeframe=bt.TimeFrame.Days,
        annualize=True,
        riskfreerate=0,
    )  # 计算夏普比率
    cerebro.addanalyzer(bt.analyzers.SharpeRatio_A, _name="_SharpeRatio_A")
    # 返回收益率时序
    cerebro.addanalyzer(bt.analyzers.TimeReturn, _name="_TimeReturn")

    # Print out the starting conditions
    print("Starting Portfolio Value: %.2f" % cerebro.broker.getvalue())
    cerebro.broker.set_coc(True)  # 设置以当日收盘价成交
    # Run over everything
    results = cerebro.run()
    strat = results[0]
    # Print out the final result
    print("Final Portfolio Value: %.2f" % cerebro.broker.getvalue())
    daily_return = pd.Series(strat.analyzers.pnl.get_analysis())
    # print(daily_return)

    # # 从返回的 result 中提取回测结果
    # strat = result[0]
    # # 返回日度收益率序列
    # # 打印评价指标
    print("AnnualReturn= ", strat.analyzers._AnnualReturn.get_analysis())
    print("SharpeRatio= ", strat.analyzers._SharpeRatio.get_analysis())
    print("DrawDown= ", strat.analyzers._DrawDown.get_analysis())
    print("_Returns= ", strat.analyzers._Returns.get_analysis())

    # Plot the result
    results = cerebro.plot(style="candle", barup="red", bardown="green", volume=True)


# def self_KDJ(data, fastk_period=9, slowk_period=3, slowd_period=3):
#     # 计算kd指标
#     high_prices = np.array([v["high"] for v in data])
#     low_prices = np.array([v["low"] for v in data])
#     close_prices = np.array([v["close"] for v in data])
#     fast_k = RSV(data, fastk_period)
#     slow_k = talib.MA(np.array(fast_k), timeperiod=slowk_period)
#     slow_d = talib.MA(slow_k, timeperiod=slowd_period)
#     indicators = {"k": slow_k, "d": slow_d, "j": 3 * slow_k - 2 * slow_d}
#     return indicators


# 定义Observer
class OrderObserver(bt.observer.Observer):
    lines = (
        "created",
        "expired",
    )
    # 做图参数设置
    plotinfo = dict(plot=True, subplot=True, plotlinelabels=True)
    # 创建工单 * 标识，过期工单 方块 标识
    plotlines = dict(
        created=dict(marker="*", markersize=8.0, color="lime", fillstyle="full"),
        expired=dict(marker="s", markersize=8.0, color="red", fillstyle="full"),
    )

    # 处理 Lines
    def next(self):
        for order in self._owner._orderspending:
            if order.data is not self.data:
                continue

            if not order.isbuy():
                continue

            # Only interested in "buy" orders, because the sell orders
            # in the strategy are Market orders and will be immediately
            # executed

            if order.status in [bt.Order.Accepted, bt.Order.Submitted]:
                self.lines.created[0] = order.created.price

            elif order.status in [bt.Order.Expired]:
                self.lines.expired[0] = order.created.price


class KDJ(bt.Indicator):
    lines = ("K", "D", "J")

    params = (
        ("period", 9),
        ("period_dfast", 3),
        ("period_dslow", 3),
    )

    plotlines = dict(
        J=dict(
            _fill_gt=("K", ("red", 0.50)),
            _fill_lt=("K", ("green", 0.50)),
        )
    )

    def __init__(self):
        # Add a KDJ indicator
        self.kd = bt.indicators.StochasticFull(
            self.data,
            period=self.p.period,
            period_dfast=self.p.period_dfast,
            period_dslow=self.p.period_dslow,
        )

        self.lines.K = self.kd.percD
        self.lines.D = self.kd.percDSlow
        self.lines.J = self.K * 3 - self.D * 2


class MyMaMacdRsi(bt.Strategy):
    params = (
        ("highperiod", 9),
        ("lowperiod", 9),
        ("kperiod", 3),
        ("dperiod", 3),
        ("period_ema5", 5),
        ("period_ema10", 10),
        ("period_ema20", 20),
        ("period_ema60", 60),
    )

    # 打印日志
    def log(self, txt, dt=None):
        dt = dt or self.data.datetime.date(0)
        print("%s, %s" % (dt, txt))

    def stop(self):
        print("stop 停止")
        self.conn.close()

    def start(self):
        print("start 开始")
        self.conn = get_mysql_conn()

    def __init__(self):
        # 用于保存订单
        self.order = None
        # 订单价格
        self.buyprice = None
        # 订单佣金
        self.buycomm = None

        # 定义变量保存所有收盘价
        self.dataopen = self.data.open
        self.datahigh = self.data.high
        self.datalow = self.data.low
        self.dataclose = self.data.close
        self.datavolume = self.data.volume
        # self.datacode = self.data.code
        #

        # 计算5日均线
        # self.ema5 = bt.talib.EMA(self.datavolume, timeperiod=self.params.period_ema5)
        # 计算10日均线
        self.ema10 = bt.talib.EMA(self.dataclose, timeperiod=self.params.period_ema10)
        # 计算20日均线
        self.ema20 = bt.talib.EMA(self.dataclose, timeperiod=self.params.period_ema20)
        # print("ema20 = ", self.ema20)
        # # 计算60日均线
        self.ema60 = bt.talib.EMA(self.dataclose, timeperiod=self.params.period_ema60)

        # N个交易日内最高价
        # self.highest = bt.indicators.Highest(
        #     self.data.high, period=self.params.highperiod
        # )
        # # N个交易日内最低价
        # self.lowest = bt.indicators.Lowest(self.data.low, period=self.params.lowperiod)

        # # 计算rsv值 RSV=(CLOSE- LOW) / (HIGH-LOW) * 100
        # # 如果被除数0 ，为None
        # self.rsv = 100 * bt.DivByZero(
        #     self.dataclose - self.lowest, self.highest - self.lowest, zero=None
        # )

        # # 计算rsv的N个周期加权平均值，即K值
        # self.K = bt.indicators.EMA(self.rsv, period=self.params.kperiod, plot=False)
        # # D值=K值 的N个周期加权平均值
        # self.D = bt.indicators.EMA(self.K, period=self.params.dperiod, plot=False)
        # # J=3*K-2*D
        # self.J = 3 * self.K - 2 * self.D

        # # 指数均线
        # bt.indicators.ExponentialMovingAverage(self.datas[0], period=21)

        # # 加权均线
        # bt.indicators.WeightedMovingAverage(self.datas[0], period=21, subplot=True)
        # #  慢速随机指数
        # bt.indicators.StochasticSlow(self.datas[0])
        # # 异同移动平均线
        # bt.indicators.MACDHisto(self.datas[0])
        # # 相对强弱指数
        # rsi = bt.indicators.RSI(self.datas[0])
        # # 平均相对强弱指数
        # bt.indicators.SmoothedMovingAverage(rsi, period=5)
        # 平均真实波动范围
        # bt.indicators.ATR(self.datas[0], plot=True)

        # Chaikin A/D Line 累积/派发线（Accumulation/Distribution Line） 简介：Marc Chaikin提出的一种平衡交易量指标
        # self.dataad = bt.talib.AD(
        #     self.datahigh, self.datalow, self.dataclose, self.datavolume
        # )

        # atr
        # self.dataatr = bt.talib.ATR(
        #     self.datahigh, self.datalow, self.dataclose, timeperiod=14
        # )

        # 平均趋向指数 简介：使用ADX指标，指标判断盘整、振荡和单边趋势。
        # self.dataadx = bt.talib.ADX(
        #     self.datahigh, self.datalow, self.dataclose, timeperiod=14
        # )

        # 投资学中意思为续航，指股票(或经济指数)持续增长的能力
        # self.datamom = bt.talib.MOM(self.dataclose)

        self.datamacd = bt.indicators.MACDHisto(self.datas[0])

        # 相对强弱指数
        self.dataris = bt.indicators.RSI(self.datas[0])
        # print("data ris = ", self.dataris)

        self.datakdj = KDJ(self.datas[0])

        # # 顺势指标
        # self.datacci = bt.talib.CCI(self.datahigh, self.datalow, self.dataclose)

        # # 资金流量指标
        # self.datamfi = bt.talib.MFI(
        #     self.datahigh, self.datalow, self.dataclose, self.datavolume
        # )

        # # 上升动向值
        # self.datamom = bt.talib.MOM(self.dataclose, timeperiod=10)

        # 定义全局变量
        self.count = 0
        for data in self.datas:
            # 转为tabib要求的数据格式
            opens = np.array(data.open.array)
            highs = np.array(data.high.array)
            lows = np.array(data.low.array)
            closes = np.array(data.close.array)
            # print(opens)
            # 数据放入self中

            # 计算十字星数据，结果为-100底部十字星，结果为100顶部十字星，0非十字星
            # res = talib.CDLDOJISTAR(opens, highs, lows, closes)
            # print("十字星数据", res)
            # self.star[data._id] = res

    def save_opt_flow(self, opt_flow):
        self.conn.insert("t_stock_trade_flow", opt_flow)
        pass

    def notify_order(self, order):
        # print("order = ", order)
        # print("order.executed = ", order.executed)
        # 等待订单提交、订单被cerebro接受
        if order.status in [order.Submitted, order.Accepted]:
            return

        # 等待订单完成
        if order.status in [order.Completed]:
            if order.isbuy():
                # self.log(
                #     "买 EXECUTED, 开仓价格: %.2f, 花费: %.2f, 手续费: %.2f"
                #     % (order.executed.price, order.executed.value, order.executed.comm)
                # )
                self.log(
                    "BUY EXECUTED, ref:%.0f，Price: %.2f, Cost: %.2f, Comm %.2f, Size: %.2f, Stock: %s"
                    % (
                        order.ref,  # 订单编号
                        order.executed.price,  # 成交价
                        order.executed.value,  # 成交额
                        order.executed.comm,  # 佣金
                        order.executed.size,  # 成交量
                        order.data._name,
                    )
                )  # 股票名称
                print(
                    "self.dataopen ",
                    self.dataopen[0],
                    "self.dataclose[0] ",
                    self.dataclose[0],
                    "ma20 ",
                    self.ema20[0],
                    "ma60 ",
                    self.ema60[0],
                    # "self.dataatr=",
                    # self.dataatr[0],
                    # "self.dataad=",
                    # self.dataad[0],
                )
                # opt_flow_data = {
                #     "id": 0,
                #     "trade_time": self.data.datetime.date(0),
                #     "code": self.data.code,
                #     "open": self.data["open"],
                #     "high": row["high"],
                #     "low": row["low"],
                #     "close": row["close"],
                #     "volume": row["volume"],
                #     "amount": row["amount"],
                #     "create_time": datetime.datetime.now(),
                #     "update_time": datetime.datetime.now(),
                # }
                # save_opt_flow(opt_flow_data)

            else:
                # self.log(
                #     "卖 EXECUTED, 卖出价格: %.2f, 收回资金: %.2f, 手续费: %.2f"
                #     % (order.executed.price, order.executed.value, order.executed.comm)
                # )
                self.log(
                    "SELL EXECUTED, ref:%.0f, Price: %.2f, Cost: %.2f, Comm %.2f, Size: %.2f, Stock: %s"
                    % (
                        order.ref,
                        order.executed.price,
                        order.executed.value,
                        order.executed.comm,
                        order.executed.size,
                        order.data._name,
                    )
                )
                self.order = None

        # 如果订单保证金不足，将不会完成，而是执行以下拒绝程序
        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log("拒绝执行，Order Canceled/Margin/Rejected = ", order.status)

        """
        买单：
        Ref: 11
        OrdType: 0
        OrdType: Buy
        Status: 1
        Status: Submitted
        Size: 10
        Price: None
        Price Limit: None
        TrailAmount: None
        TrailPercent: None
        ExecType: 0
        ExecType: Market
        CommInfo: None
        End of Session: 738732.9999999999
        Info: AutoOrderedDict()
        Broker: None
        Alive: True
        
        卖单：
        Ref: 12
        OrdType: 1
        OrdType: Sell
        Status: 1
        Status: Submitted
        Size: -10
        Price: None
        Price Limit: None
        TrailAmount: None
        TrailPercent: None
        ExecType: 0
        ExecType: Market
        CommInfo: None
        End of Session: 738740.9999999999
        Info: AutoOrderedDict()
        Broker: None
        Alive: True
        
        """

    def notify_trade(self, trade):
        if not trade.isclosed:
            return

        # self.log(
        #     "开始交易 PROFIT, GROSS %.2f, NET %.2f" % (trade.pnl, trade.pnlcomm)
        # )  # pnl：盈利  pnlcomm：手续费
        self.log(f"策略收益：毛收益 {trade.pnl:.2f}, 净收益 {trade.pnlcomm:.2f}")

    # 策略逻辑实现
    def next(self):

        print("当前可用资金", self.broker.getcash())
        print("当前总资产", self.broker.getvalue())
        print("当前持仓量", self.broker.getposition(self.data).size)
        print("当前持仓成本", self.broker.getposition(self.data).price)
        # 也可以直接获取持仓
        print("当前持仓量", self.getposition(self.data).size)
        print("当前持仓成本", self.getposition(self.data).price)
        print("self.datakdj = ", self.datakdj)
        # 跌破10日均线
        cur_down_cross_10_line = self.dataclose[0] < self.ema10[0]
        # 10日均线小于20日
        cur_ema_10_lt_20 = self.ema10[0] < self.ema20[0]
        cur_ema_20_gt_pre_20 = self.ema20[0] > self.ema20[-1]
        # 20日均线大于60日
        cur_ema_20_gt_60 = self.ema20[0] > self.ema60[0]

        # macd 金叉，macd零轴之上，macd上升趋势
        cur_macd_jincha = self.datamacd.histo[0] > 0
        cur_macd_gt_0 = self.datamacd.macd[0] >= 0
        cur_macd_up = self.datamacd.macd[0] > self.datamacd.macd[-1]

        # rsi 大于50，rsi上升趋势
        cur_rsi_gt_50 = self.dataris[0] > 50
        # K线当日收阳线，大于昨日上升趋势，
        cur_k_line_up = self.dataclose > self.dataopen

        cur_volume_up = self.datavolume[0] > self.datavolume[-1]

        cur_k_qushi = (
            self.dataclose[0] > self.dataclose[-1]
            and self.datahigh[0] > self.datahigh[-1]
            and self.datalow[0] > self.datalow[-1]
        )

        cur_jiang = self.dataclose[0] < self.datalow[-1]

        is_kai_cang = (
            cur_macd_jincha
            and cur_macd_up
            and cur_rsi_gt_50
            and cur_k_line_up
            and cur_k_qushi
            and cur_volume_up
        )

        # 当今天的20日均线大于60日均线并且昨天的20日均线小于60日均线，则进入市场（买）
        # if self.ema20[0] > self.ema60[0] and self.ema20[-1] < self.ema60[-1]:
        # if (
        #     cur_ema_20_gt_60
        #     and cur_macd_jincha
        #     and cur_macd_gt_0
        #     and cur_macd_up
        #     and cur_rsi_gt_50
        #     and cur_k_line_up
        #     and cur_k_qushi
        # ):
        if is_kai_cang:
            # print("可以考虑买入，self.order = ", self.order)
            # 判断订单是否完成，完成则为None，否则为订单信息
            if not self.order is None:
                # self.log("已经存在订单了")
                return
            post_total = self.dataclose[0] * 100
            print("需要花费：", post_total, "元")

            # 若上一个订单处理完成，可继续执行买入操作
            self.order = self.buy()
            # print("已经买入，self.order = ", self.order)
            print("已经买入 close_price=", self.dataclose[0], " ok!!!")

        # 当今天的20日均线小于60日均线并且昨天的20日均线大于60日均线，则退出市场（卖）
        # elif (
        #     cur_ema_10_lt_20
        #     or not cur_ema_20_gt_60
        #     or not cur_macd_jincha
        #     or not cur_macd_gt_0
        #     or not cur_rsi_gt_50
        #     # or cur_jiang
        # ):
        elif cur_down_cross_10_line:

            # print("需要卖出，self.order = ", self.order)
            if self.order is None:
                # self.log("没有订单需要卖出")
                return
            # 卖出
            self.order = self.sell()
            # print("已经卖出，self.order = ", self.order)
            print("已经卖出，close_price=", self.dataclose[0])

        # for data in self.datas:
        #     pos = self.getposition(data).size
        # 函数出现100就代表

        # 自增处理
        self.count = self.count + 1


if __name__ == "__main__":
    code = "000333"
    # code = "002594"
    get_data_and_read(code)
    backtrader_test_by_data(code)
